(* $Id: ergoPrinter.ml 84 2007-10-17 08:39:00Z yann.regisgianas $ *)

(* Pangolin, a functional programming language for correct programs.
   Copyright (C) 2007, Yann Régis-Gianas, François Pottier.
   Contact: yann.regisgianas@gmail.com

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>. *)

module B = Buffer
open Pprint
open PprintExt
open PFOLSyntax.Raw

let id_text x = 
  text (Demangle.str x)

let kw x = 
  text x

let sym x =
  text x

let pid x = 
  text (PIdentifier.as_string x)

let tid x = 
  text (PIdentifier.as_string x)

let indent t = 
  hang 2 (group t)
    
let fresh_counter = ref 0

let rec name_from_int i : string =
  if i < 26 then
    String.make 1 (Char.chr (0x61 + i))
  else
    String.concat "_" [ name_from_int (i / 26 - 1) ; name_from_int (i mod 26) ]

let init_fresh () = 
  fresh_counter := 0

let fresh () = 
  incr fresh_counter;
  !fresh_counter

let fresh_type_id _ =
  tid (PIdentifier.mk_type_id (Position.dummy, 
			      "'" ^ name_from_int (fresh ())))

module IntSet = Set.Make (struct type t = int
				 let compare = compare 
			  end)

let used_tuples = ref IntSet.empty 

let tuple_name i = 
  Printf.sprintf "tuple%d" i

let tuple i = 
  used_tuples := IntSet.add i !used_tuples;
  tuple_name i

let tuple_type i = 
  let vs = Misc.repeat i fresh_type_id in
  let d = comma_sep vs in
    d ++ sym "->" ++ paren d ++ text (tuple i)
      
let tuple_type_declaration i = 
  let vs = Misc.repeat i fresh_type_id in
    kw "type" ++ paren (comma_sep vs) ++ text (tuple_name i)

let declare_tuple i ds = 
  ((tuple_type_declaration i $$
      kw "logic " ++ text (tuple_name i) ++ sym ":" ++ tuple_type i)
  ) :: ds

let declare_tuples () = 
  vcat (IntSet.fold declare_tuple !used_tuples [])

let fresh_axiom_id _ =
  pid (PIdentifier.mk_predicate_id (Position.dummy, 
				    (Printf.sprintf "property_%d" (fresh ()))))

let mark_as_universal rs v = 
  PIdentifier.Set.add v rs

let mark_as_universals rs vs = 
  List.fold_left mark_as_universal rs vs

let is_universal v rs =
  (PIdentifier.Set.mem v rs)

let empty_rs = PIdentifier.Set.empty

let universals_as_string s = 
  PIdentifier.Set.fold (fun x s -> PIdentifier.as_string x ^ " " ^ s) s ""

let rec fl_quantified q quantifier bs (ts : trigger list) body t = 
  match bs with
    | [] -> 
	body t

    | [ v ] -> 
	group (kw q ++ quantifier v ^^ triggers ts ^^ sym "." 
		 $$ group (body t))
    | v :: bs ->
	group (kw q ++ quantifier v ^^ sym "." 
		 $$ group (fl_quantified q quantifier bs ts body t))

and ex_quantified q quantifier ts body t = 
  match ts with
    | [] -> body t
    | v :: ts ->
	group (kw q ++ quantifier v ^^ sym "." 
		 $$ group (ex_quantified q quantifier ts body t))

and triggers ts = 
  if ts = [] then empty
  else lbracket ++ comma_sep (List.map tid ts) ++ rbracket

let rec predicate rs p = 
  match p with
    | ForallTys (vs, p) ->
	predicate (mark_as_universals rs vs) p
	  
    | Forall (bs, ts, p) ->
	fl_quantified "forall" (binding rs) bs ts (predicate rs) p
	  
    | Exists (bs, p) ->
	ex_quantified "exists" (binding rs) bs (predicate rs) p
	  
    | LogicBinOpApp (op, lhs, rhs) ->
	group 
	  (paren
	     (paren (predicate rs lhs) $$ binop op ++ paren (predicate rs rhs)))

    | BinOpApp (op, lhs, rhs) ->
	group 
	  (paren
	     (paren (term lhs) $$ binop op ++ paren (term rhs)))
	  
    | Neg p ->
	paren (kw "not" ++ predicate rs p)
	 
    | PApp (p, []) ->
	id p
 
    | PApp (p, ts) ->
	id p ++ tupled' term ts
	  
    | PEq (lhs, rhs) ->
	paren (term lhs ++ sym "=" ++ term rhs)
	  
    | PTrue ->
	kw "true"
	  
    | PFalse ->
	kw "false"
	  
    | Pre (f, t, u) ->
	kw "pre" ++ tupled [id f; term t; term u ]
	  
    | Post (f, t1, t2, t3) ->
	kw "post" ++ tupled [id f; term t1; term t2; term t3 ]

    | Reify t ->
	kw "reify" ++ paren (term t)

and ipredicate rs p = 
  indent (predicate rs p)

and binop = function
  | Conj -> kw "and"
  | Disj -> kw "or"
  | Imply -> kw "->"
  | Equiv -> kw "<->"
  | LessThan -> kw "<"
  | LessEqualThan -> kw "<="
  | GreaterThan -> kw ">"
  | GreaterEqualThan -> kw ">="
	
and bindings rs bs = 
    tupled' (binding rs) bs

and binding rs (x, ty) =
  id x ++ sym ":" ++ term_type rs ty

and term = function
  | Var x -> 
      id x

  | App (f, ts) -> 
      id f ++ tupled' term ts

  | Prod [] ->
      text "void"

  | Prod [ t ] ->
      term t

  | Prod ts -> 
      text (tuple (List.length ts)) ++ tupled' term ts

  | Prim p ->
      primitive p

  | PrimApp (p, [ t1; t2 ]) ->
      paren (term t1) ++ term_binop p ++ paren (term t2)

  | PredObjectApp (p, ts) ->
      kw "apply" ++ tupled [ term p; term (Prod ts) ]

  | Annot (p, _) ->
      term p

  | _ -> 
      assert false

and term_binop = function
  | PAdd -> 
      sym "+"

  | PSub ->
      sym "-"

  | PMult ->
      sym "*"

  | PDiv ->
      sym "/"

  | _ -> 
      assert false


and primitive = function
  | PInt x ->
      text (string_of_int x)

  | PTrueBool ->
      text "true"

  | PFalseBool ->
      text "false"

  | PLessThanObj ->
      text "less_than_object"

  | PGreaterThanObj ->
      text "greater_than_object"

  | PLessEqualThanObj ->
      text "less_equal_than_object"

  | PGreaterEqualThanObj ->
      text "greater_equal_than_object"

  | PEqObj ->
      text "equal_obj"

  | _ -> 
      assert false

and term_type rs = function
  | TVar x -> 
      if is_universal x rs then sym "'" ^^ tid x else tid x

  | TApp (x, []) ->
      if is_universal x rs then sym "'" ^^ tid x else tid x

  | TApp (tycon, [t]) ->
      term_type rs t ++ tid tycon

  | TApp (tycon, ts) ->
      tupled' (term_type rs) ts ++ tid tycon

  | TCArrow (ins, outs) ->
      paren (term_type rs ins ++ sym "," ++ term_type rs outs) ++ kw "fun"

  | TPArrow (ins, outs) ->
      paren (term_type rs ins ++ sym "," ++ term_type rs outs) ++ kw "pred"

  | TArrow (TProd ins, outs) ->
      comma_sep' (term_type rs) ins ++ sym "->" ++ term_type rs outs

  | TArrow (ins, outs) ->
      term_type rs ins ++ sym "->" ++ term_type rs outs

  | TProd [] ->
      kw "unit"

  | TProd [ t ] ->
      term_type rs t

  | TProd ts ->
      paren (comma_sep' (term_type rs) ts) 
      ++ text (tuple (List.length ts))

  | TPrimitive p ->
      primitive_type p

  | TProp ->
      kw "object_prop"

  | _ ->
      assert false

and primitive_type = function
  | TInt -> kw "int"
  | TBool -> kw "bool"
  | TUnit -> kw "unit"

and arity k = 
  if k = 0 then
    empty
  else 
    tupled (Misc.repeat k fresh_type_id)

and type_scheme_from_bindings rs bs = 
  tupled' (term_type rs) (List.map (fun (_, ty) -> ty) bs)

and declaration = function
  | DType (t, k) ->
      kw "type" ++ arity k ++ tid t

  | DVal (v, ts) -> 
      kw "logic" ++ id v ++ sym ":" ++ type_scheme ts
	
  | DPred (p, PredDec s) ->
      kw "logic" ++ id p ++ sym ":" ++ predicate_type_scheme s
	
  | DPred (p, PredDef (vs, bs, def)) ->
      let rs = mark_as_universals empty_rs vs in
	text "(* PREDICATE DECLARATION. *)"
	  $$ kw "predicate" ++ id p ++ bindings rs bs 
	++ sym "=" $ ipredicate rs def
	  
  | DAxiom (vs, f) ->
      let rs = mark_as_universals empty_rs vs in
	kw "axiom" ++ fresh_axiom_id () ++ sym ":" $ ipredicate rs f

and type_scheme (vs, ty) =
  term_type (mark_as_universals empty_rs vs) ty

and predicate_type_scheme (vs, tys) =
  let rs = mark_as_universals empty_rs vs in
  match tys with
    | [ ty ] -> term_type rs ty ++ sym "->" ++ kw "prop"
    | tys -> comma_sep' (term_type rs) tys ++ sym "->" ++ kw "prop"

and type_parameters ts =
    empty

and type_parameter t = 
    empty 

and id t = 
  id_text (PIdentifier.as_string t)
    
let rec query = function
  | Goal p -> 
      let g = predicate empty_rs p in
	group (kw "goal" ++ text "pangolin" ++ sym ":" ++ g), g

  | ConsDecl (d, q) ->
      let q, g = query q in
	(indent (declaration d) $$ q), g

let as_string d = 
  let b = B.create 13 in
    Buffer.pretty 0.6 120 b d;
    B.contents b  

let int_binop_type = 
  text "((int, int)" ++ text (tuple 2) ++ text ", object_prop) pred"

let poly_eq_type = 
  text "(('a, 'a)" ++ text (tuple 2) ++ text ", object_prop) pred"

let prologue = 
  empty
    $$ kw "type" ++ text "('a, 'b) fun" 
    $$ kw "type" ++ text "('a, 'b) pred" 
    $$ kw "type" ++ text "object_prop"
    $$ kw "logic" ++ text "apply : ('a, 'b) pred, 'a -> 'b"
    $$ kw "logic" ++ text "reify : object_prop -> prop"
    $$ kw "logic" ++ text "pre: ('a, 'b) fun, 'c, 'a -> prop"
    $$ kw "logic" ++ text "post: ('a, 'b) fun, 'c, 'a, 'b -> prop"
    $$ kw "logic" ++ text "less_than_object:" ++ int_binop_type
    $$ kw "axiom less_reify:" ++ text 
    "forall x, y: int. reify (apply (less_than_object, tuple2 (x, y))) <-> x < y"
    $$ kw "logic" ++ text "greater_than_object:" ++ int_binop_type
    $$ kw "axiom greater_reify:" ++ text 
    "forall x, y: int. reify (apply (greater_than_object, tuple2 (x, y))) <-> x > y"
    $$ kw "logic" ++ text "less_equal_than_object:" ++ int_binop_type
    $$ kw "axiom less_equal_reify:" ++ text 
    "forall x, y: int. reify (apply (less_equal_than_object, tuple2 (x, y))) \
     <-> x <= y"
    $$ kw "logic" ++ text "greater_equal_than_object:" ++ int_binop_type
    $$ kw "axiom greater_equal_reify:" ++ text 
    "forall x, y: int. reify (apply (greater_equal_than_object, tuple2 (x, y))) \
     <-> x >= y"
    $$ kw "logic" ++ text "equal_object:" ++ poly_eq_type
    $$ kw "axiom equal_reify:" ++ text 
    "forall x, y: 'a. reify (apply (equal_object, tuple2 (x, y))) <-> x = y"

let print p = 
  init_fresh ();
  let d, g = query p in
    (as_string (declare_tuples () $$ prologue $$ d), as_string g)


